#include "typeHelpers.hpp"
#include "unifiedMathFunctions.hpp"
#include "eigenHelpers.hpp"
#include "eigenConcat.hpp"
#include "exceptions.hpp"
#include "macros.hpp"



namespace hms::swe::hllc
{

using WaveRegion = hms::hllc::WaveRegion;

template<typename T>
void assert1xN_or_arithmetic(
	[[maybe_unused]] const T& t,
	[[maybe_unused]] Index nCols
){
	using U = remove_qualifiers<T>;
	if constexpr ( isDenseType<U>() ){
		static_assert( U::RowsAtCompileTime == 1 );
		assert( t.cols() == nCols );
	} else {
		static_assert( std::is_arithmetic_v<U> );
	}
}


/** @brief Returns the HLLC wave speed for the specified side (left or right),
 * in the case that that side is dry, but the opposite one is not.
 * 
 * Implements Simons (2019), Eq. 3.54 and 3.55, first branch each
 * @see hms::hllc::wavespeed() */
template<WaveRegion side, typename V, typename C>
auto wavespeed_dry(
	const V& vNormalOther,
	const C& cwOther
) -> decltype(auto)
{
	static constexpr int sign { side == WaveRegion::left ? -1 : 1 };
	static_assert( side == WaveRegion::left || side == WaveRegion::right );
	return vNormalOther + sign * 2 * cwOther;
}


/** @brief Returns the HLLC wave speed for the specified side (left or right),
 * in the case that that side is not dry.
 * 
 * @param vNormalSide Velocity normal to the edge, flux' side.
 * @param vNormalStar Velocity normal to the edge in the star region.
 * @param cwSide The wavespeed c_W, not the HLLC one, on the flux' side.
 * @param cwStar The wavespeed c_W*, not the HLLC one.
 * 
 * Implements Simons (2019), Eq. 3.54 and 3.55, second branch each.
 * @see wavespeed_dry() */
template<WaveRegion side,
	typename VSide, typename VStar, typename CSide, typename CStar>
auto wavespeed(
	const VSide& vNormalSide, const VStar& vNormalStar,
	const CSide&      cwSide, const CStar&      cwStar
) -> decltype(auto)
{
    static_assert( side == WaveRegion::left || side == WaveRegion::right );

	if constexpr ( side == WaveRegion::left )
		return min(vNormalSide - cwSide, vNormalStar - cwStar);
	else
		return max(vNormalSide + cwSide, vNormalStar + cwStar);
}


/**
 * @brief Computes the hllc wave speeds. Checks whether water depth is zero
 * and calls wavespeed_dry accordingly.
 * 
 * @tparam side 
 * @param dSide 
 * @param vNormalSide 
 * @param vNormalOther 
 * @param cwSide 
 * @param cwOther 
 * @param cwStar 
 * @return decltype(auto) 
 */
template<WaveRegion side,
	typename D, typename VSide, typename VStar, typename CSide, typename CStar>
auto wavespeed(
	const D& dSide,
	const VSide& vNormalSide, const VSide& vNormalOther, const VStar& vNormalStar,
	const CSide&      cwSide, const CSide&      cwOther, const CStar&      cwStar
) -> decltype(auto)
{
	assert( notNegative(dSide) );

	return HMS_SELECT( dSide == 0,
		wavespeed_dry<side>(vNormalOther, cwOther),
		wavespeed    <side>(vNormalSide, vNormalStar, cwSide, cwStar)
	);
}


/**
 * @brief According to Simons (2019), Eq. 3.50,
 * the wavespeed in the star region S* is used to decide whether 
 * f_*L or f_*R is used to calculate the flux. However, when SL < 0 && SR > 0,
 * then the only condition that remains is S* >= 0 or S* < 0.
 * This can be done via a sign comparison between numerator and denominator.
 * Same speed.
 * With this variant, no division by zero occurs, but is that better?
 * 
 * @param dL water depth left (1xN)
 * @param dR water depth right (1xN)
 * @param uL normal velocity left (1xN)
 * @param uR normal velocity right (1xN)
 * @param sL hllc wavespeed left (1xN)
 * @param sR hllc wavespeed right (1xN)
 * @return decltype(auto) 
 */
template<typename D, typename V, typename Sl, typename Sr>
auto useFluxLeft(
	const D& dL, const D& dR,
	const V& vL, const V& vR,
	const Sl& sL, const Sr& sR
) -> decltype(auto)
{
	assert( notNegative(dL) );
	assert( notNegative(dR) );
    // auto hL { dL * ( vL - SL ) };
    // auto hR { dR * ( vR - SR ) };
	// return sign(SL * hR - SR * hL) == sign(hR - hL);
    return
		sign( sL * dR * (vR - sR) - sR * dL * (vL - sL) ) ==
		sign(      dR * (vR - sR) -      dL * (vL - sL) );
}


template<typename ... Ts>
decltype(auto) toFluxComponents( Ts&& ... args ){
	if constexpr (
		( std::is_base_of_v<
			Eigen::DenseBase<remove_qualifiers<Ts>>, remove_qualifiers<Ts>
		> && ... )
	){
		return concat_vertical( std::forward<Ts>(args) ... );
	} else {
		static_assert( (std::is_arithmetic_v<remove_qualifiers<Ts>> && ...) );
		return ArrayNs<sizeof...(Ts)>{ std::forward<Ts>(args) ... };
	}
}


/**
 * @brief General form of a flux for non-orthogonal meshes as
 *            |           f1              |
 * F(f1,f2) = | nx * f2 - ny * f1 * vPara |
 *            | ny * f2 + nx * f1 * vPara |
 * 
 * with nx and ny being the normal vector components and vPara being the 
 * velocity parallel to the edge.
 * 
 * @param f1 first component of a flux vector, either HLL flux (f*) or
 * the flux in the left/right regions (f_L, f_R)
 * @param f2 second component of said flux vector
 * @param vPara velocities parallel to an edge
 * @param normal normal vectors
 * @return decltype(auto) 
 * 
 * f1,f2 are components of some vector. For the left and right region HLLC
 * fluxes, this would be
 * f1 = vNormal * d and
 * f2 = vNormal^2 * d + 1/2 g * d^2
 * For the star region, this is f*1 and f*2, the components of the HLL flux.
 * Implements Simons (2018), Eq. 3.63-3.65.
 */
template<typename F1, typename F2, typename V, typename N>
auto flux(
	const F1& f1,
	const F2& f2,
	const V& vPara,
	const N& normal
) -> decltype(auto)
{
	static_assert( N::RowsAtCompileTime == 2 );
	assert1xN_or_arithmetic( f1, normal.cols() );
	assert1xN_or_arithmetic( f2, normal.cols() );
	assert1xN_or_arithmetic( vPara, normal.cols() );

	auto n = [&](Index row) -> decltype(auto) {
		if constexpr ( isDenseType<F1>() ){
			return normal.row(row).array();
		} else {
			static_assert( N::ColsAtCompileTime == 1 );
			return normal[row];
		}
	};

	return toFluxComponents(
		f1,
		n(0) * f2 - n(1) * f1 * vPara,
		n(1) * f2 + n(0) * f1 * vPara
	);
}

/**
 * @brief General form of a flux on orthogonal meshes, where normal vectors
 * are either [1,0] or [0,1] and thus several terms are eliminated.
 * 
 * @param f 
 * @param vPara 
 * @return decltype(auto)
 * 
 *            |    f1      |
 * F(f1,f2) = |    f2      |
 *            | f1 * vPara |
 */
template<typename F, typename V>
auto flux_ortho(
	const F& f,
	const V& vPara
) -> decltype(auto)
{
	static_assert( F::RowsAtCompileTime == 2 );

	assert1xN_or_arithmetic( vPara, f.cols() );

	return concat_vertical(
		f, /* first two rows are just f1 and f2 */
		f.row(0) * vPara /* third row is f1 * vPara */
	);
}


/**
 * @brief Flux in the outer regions (left/right).
 * 
 * @param d 
 * @param vNormal 
 * @param vPara 
 * @param normal 
 * @return decltype(auto) 
 */
template<typename D, typename Vn, typename Vp, typename N>
auto flux_outer(
	const D& d,
	const Vn& vNormal,
	const Vp& vPara,
	const N& normal
) -> decltype(auto)
{
	static_assert( N::RowsAtCompileTime == 2 );

	#ifndef NDEBUG
	Index nCols { normal.cols() };

	assert1xN_or_arithmetic( d, nCols );
	assert1xN_or_arithmetic( vPara, nCols );
	assert1xN_or_arithmetic( vNormal, nCols );
	#endif

	return flux(
		massFlux(d, vNormal),
		momentumFlux_normal(d, vNormal),
		vPara,
		normal
	);
}

template<
	typename D,
	typename Vn1, typename Vn2,
	typename Vp, typename N,
	typename S1, typename S2>
auto flux_star(
	const D& dSide,
	const D& dOther,
	const Vn1& vNormalSide,
	const Vn2& vNormalOther,
	const Vp& vPara,
	const N& normal,
	const S1& sSide,
	const S2& sOther
) -> decltype(auto)
{
	static_assert( N::RowsAtCompileTime == 2 );

	#ifndef NDEBUG
	Index nCols { normal.cols() };

	assert1xN_or_arithmetic( dSide , nCols );
	assert1xN_or_arithmetic( dOther, nCols );
	assert1xN_or_arithmetic( vNormalSide , nCols );
	assert1xN_or_arithmetic( vNormalOther, nCols );
	assert1xN_or_arithmetic( vPara, nCols );
	assert1xN_or_arithmetic( sSide , nCols );
	assert1xN_or_arithmetic( sOther, nCols );
	#endif

	return flux(
		hllFlux(
			sSide, sOther,
			massFlux(dSide, vNormalSide), massFlux(dOther, vNormalOther),
			dSide, dOther
		),
		hllFlux(
			sSide, sOther,
			momentumFlux_normal(dSide, vNormalSide),
			momentumFlux_normal(dOther, vNormalOther),
			massFlux(dSide, vNormalSide), massFlux(dOther, vNormalOther)
		),
		vPara,
		normal
	);
}

/**
 * @brief Returns the flux in either the outer left or outer right region
 * (FL or FR), depending on the left HLLC wave speed SL. Caution: Only use this,
 * if you ascertained that it's an outer flux,
 * i.e. that either SL >= 0 or SR <= 0! Otherwise, use flux_star_LR.
 * 
 * @tparam D 
 * @tparam Vn 
 * @tparam Vp 
 * @tparam N 
 * @tparam S 
 * @param dL 
 * @param dR 
 * @param vNormalL 
 * @param vNormalR 
 * @param vParaL 
 * @param vParaR 
 * @param normal 
 * @param sL 
 * @return decltype(auto) 
 */
template<
	typename D,
	typename Vnl, typename Vnr,
	typename Vpl, typename Vpr,
	typename N,
	typename S>
auto flux_outer_LR(
	const D& dL,
	const D& dR,
	const Vnl& vNormalL,
	const Vnr& vNormalR,
	const Vpl& vParaL,
	const Vpr& vParaR,
	const N& normal,
	const S& sL
) -> decltype(auto)
{
	static constexpr Index nRows {3};
	return HMS_SELECT(
		repl<nRows>( sL >= 0 ),
		flux_outer(dL, vNormalL, vParaL, normal),
		flux_outer(dR, vNormalR, vParaR, normal)
	);
}

/**
 * @brief Returns flux in either the left or right star region (F*L or F*R),
 * depending on the middle wave speed S*.
 * 
 * @tparam D 
 * @tparam Vn 
 * @tparam Vp 
 * @tparam N 
 * @tparam S 
 * @param dL 
 * @param dR 
 * @param vNormalL 
 * @param vNormalR 
 * @param vParaL 
 * @param vParaR 
 * @param normal 
 * @param sL 
 * @param sR 
 * @return decltype(auto) 
 */
template<
	typename D,
	typename Vnl, typename Vnr,
	typename Vpl, typename Vpr,
	typename N,
	typename Sl, typename Sr>
auto flux_star_LR(
	const D& dL,
	const D& dR,
	const Vnl& vNormalL,
	const Vnr& vNormalR,
	const Vpl& vParaL,
	const Vpr& vParaR,
	const N& normal,
	const Sl& sL,
	const Sr& sR
) -> decltype(auto)
{
	/* number of flux components */
	static constexpr Index nRows {3};
	return HMS_SELECT(
		repl<nRows>( useFluxLeft(dL, dR, vNormalL, vNormalR, sL, sR) ),
		flux_star(dL, dR, vNormalL, vNormalR, vParaL, normal, sL, sR),
		flux_star(dR, dL, vNormalR, vNormalL, vParaR, normal, sR, sL)
	);
}

/**
 * @brief Uses reconstructed edge values to carry out the HLLC scheme for
 * advective flux computation. The parameter "flux" is an out-parameter, which
 * gets overwritten.
 * 
 * 
 * @tparam SolverType required for minWaterDepth.
 * @tparam F 
 * @tparam D 
 * @tparam Vn 
 * @tparam Vp 
 * @tparam N 
 * @param flux 
 * @param dL 
 * @param dR 
 * @param vNormalL 
 * @param vNormalR 
 * @param vParaL 
 * @param vParaR 
 * @param normal 
 */
template<typename SolverType,
	typename F, typename D,
	typename Vnl, typename Vnr,
	typename Vpl, typename Vpr,
	typename N>
void hllcScheme(
	F& flux,
	const D& dL, const D& dR,
	const Vnl& vNormalL, const Vnr& vNormalR,
	const Vpl& vParaL, const Vpr& vParaR,
	const N& normal
){
	static_assert(F::RowsAtCompileTime == 3);
	static_assert(N::RowsAtCompileTime == 2);
	assert( normal.cols() == flux.cols() );
	assert1xN_or_arithmetic( dL, flux.cols() );
	assert1xN_or_arithmetic( dR, flux.cols() );
	assert1xN_or_arithmetic( vNormalL, flux.cols() );
	assert1xN_or_arithmetic( vNormalR, flux.cols() );
	assert1xN_or_arithmetic( vParaL, flux.cols() );
	assert1xN_or_arithmetic( vParaR, flux.cols() );

	/* wave speeds */
	auto cwL { swe::wavespeed(dL) };
	auto cwR { swe::wavespeed(dR) };

	/* water depth and wave speed in the star region */
	auto  vStar {  vNormalStar(vNormalL, vNormalR, cwL, cwR) };
	auto cwStar { hllc::cwStar(vNormalL, vNormalR, cwL, cwR) };

	using Side = hllc::WaveRegion;
	auto sL { wavespeed<Side::left>(
		dL, vNormalL, vNormalR, vStar, cwL, cwR, cwStar
	) };
	auto sR { wavespeed<Side::right>(
		dR, vNormalR, vNormalL, vStar, cwR, cwL, cwStar
	) };

	/* number of flux components */
	static constexpr Index nRows { F::RowsAtCompileTime };
	static constexpr auto dryDepth { SolverType::minWaterDepth };

	auto zero = [&](){
		if constexpr ( isDenseType<D>() )
			return 0;
		else
			return ArrayNs<nRows>::Zero();
	};

	/* compute fluxes using hllc waves */
	flux = HMS_SELECT( repl<nRows>( dL < dryDepth && dR < dryDepth ), zero(),
		HMS_SELECT( repl<nRows>( sL < 0 && sR > 0 ),
		flux_star_LR (dL,dR, vNormalL,vNormalR, vParaL,vParaR, normal, sL,sR),
		flux_outer_LR(dL,dR, vNormalL,vNormalR, vParaL,vParaR, normal, sL)
	) );

// 	static constexpr auto dryDepth { SolverType::minWaterDepth };

// 	/* set fluxes to zero after the fact, if edge is dry */
// 	flux = HMS_SELECT(
// 		repl<nRows>( dL < dryDepth && dR < dryDepth ),
// 		0,
// 		flux
// 	);
}





} // namespace hms::swe::hllc
